from util.box_ops import box_cxcywh_to_xyxy
from .postprocesshoi import PostProcessHOI as posthoi
import torch
from torch import nn
import torch.nn.functional as F
import copy

from util.box_ops import box_cxcywh_to_xyxy
import numpy as np


class PostProcessHOI(posthoi):
    def __init__(self, args, mode='train', correct_mat = None, trans_type='detr'):
        super(PostProcessHOI, self).__init__(args,mode, correct_mat, trans_type)
        self.max_bin_num = 50

    @torch.no_grad()
    def forward(self, outputs, target_sizes):
        sub_boxes, obj_boxes, sub_scores, obj_scores, obj_labels, verb_scores  = \
            self.get_outputs(outputs, target_sizes)
    
        results = []
        for per_img in range(len(obj_scores)):
            os, ss, ol, vs, sb, ob =  \
                obj_scores[per_img], sub_scores[per_img], obj_labels[per_img], verb_scores[per_img], sub_boxes[per_img], obj_boxes[per_img]
            sl = torch.full_like(ol, self.subject_category_id)#0
            l = torch.cat((sl, ol))
            b = torch.cat((sb, ob))        
            s = torch.cat((ss, os))

            # hoi_scores = vs * ((os.unsqueeze(1)>0.) ** self.HO_scale)  * ((ss.unsqueeze(1)>0.) ** self.HO_scale)      
            hoi_scores = vs #* ss.unsqueeze(1) * os.unsqueeze(1)
            
            ids = torch.arange(b.shape[0])
            if self.mode=='train':
                results.append({'boxes':  np.clip(b.to('cpu'), 0, 1333),
                                'labels': l.to('cpu'),  
                                'score':  s.to('cpu')})
                results[-1].update({'sub_ids':     ids[:ids.shape[0] // 2],
                                    'obj_ids':     ids[ids.shape[0] // 2:],
                                    'verb_scores': hoi_scores.to('cpu'), 
                                    'obj_score':   os.unsqueeze(1).to('cpu'),
                                    'sub_score':   ss.unsqueeze(1).to('cpu'),})

            elif self.mode=='test':
                bboxes = [{'bbox':          np.clip(bbox, 0, 1333), 
                           'category_id':   label, 
                           'score':         score} 
                            for bbox, label, score in zip(b.to('cpu').numpy(), 
                                                          l.to('cpu').numpy(), 
                                                          s.to('cpu').numpy())]
                object_labels = ol.view(-1, 1).expand(-1, hoi_scores.shape[1])
                verb_labels = torch.arange(hoi_scores.shape[1], device=self.correct_mat.device).view(1, -1).expand(hoi_scores.shape[0], -1).to('cpu').numpy()
                subject_ids = np.tile(ids[:ids.shape[0] // 2].to('cpu').numpy(), (hoi_scores.shape[1], 1)).T
                object_ids = np.tile(ids[ids.shape[0] // 2:].to('cpu').numpy(), (hoi_scores.shape[1], 1)).T
                #   2900, 
                hoi_scores = hoi_scores.to('cpu').numpy().ravel()
                verb_labels = verb_labels.ravel()
                subject_ids = subject_ids.ravel()
                object_ids = object_ids.ravel()
                #   语义级别过滤矩阵
                masks = self.correct_mat[verb_labels.reshape(-1), 
                                         object_labels.reshape(-1)].view(hoi_scores.shape).to('cpu').numpy()
                hoi_scores *= masks
                hoi_scores[hoi_scores<self.query_thres] = 0.
                new_hoi_scores = hoi_scores[hoi_scores>0.] 
                verb_labels = verb_labels[hoi_scores>0.]
                object_ids =  object_ids[hoi_scores>0.]
                subject_ids = subject_ids[hoi_scores>0.]
                
                hois = [{'subject_id':  subject_id, 
                         'object_id':   object_id, 
                         'category_id': category_id, 
                         'score': hoi_scores} for subject_id, object_id, category_id, hoi_scores \
                             in zip(subject_ids, object_ids, verb_labels, new_hoi_scores)]
                # hois.sort(key=lambda k: (k.get('score', 0)), reverse=True)
                # hois = hois[:self.max_hois]#  取前100结果
                current_result = {'predictions': bboxes, 'hoi_prediction': hois}
                if self.use_nms_filter:
                    current_result = self.my_triplet_nms_filter(current_result)
                current_result['hoi_prediction'].sort(key=lambda k: (k.get('score', 0)), reverse=True)
                current_result['hoi_prediction'] = current_result['hoi_prediction'][:self.max_hois]
                results.append(current_result)

        return results

    @torch.no_grad()
    def get_outputs(self, outputs, target_sizes):
        assert target_sizes.shape[1] == 2
        # obj_scores_list = []
        # sub_scores_list = []
        # obj_labels_list = []
        # sub_boxes_list = []
        # obj_boxes_list = []
        # verb_scores_list = []
        # for i in range(len(outputs['pred_obj_logits_cascade'])):
        #     obj_scores, obj_labels, obj_topk_boxes = self.get_obj_cls(outputs['pred_obj_logits_cascade'][i], target_sizes)
        #     sub_scores = self.get_sub_cls(outputs['pred_sub_logits_cascade'][i], target_sizes)
        #     obj_boxes = self.get_obj_bbox(outputs['pred_obj_boxes_cascade'][i], target_sizes, obj_topk_boxes)
        #     sub_boxes = self.get_sub_bbox(outputs['pred_sub_boxes_cascade'][i], target_sizes, obj_topk_boxes)
        #     verb_scores = self.get_verb_score(outputs['pred_verb_logits_cascade'][i], target_sizes, obj_topk_boxes)
        #     verb_scores_list.append(verb_scores)
        #     obj_boxes_list.append(obj_boxes)
        #     sub_boxes_list.append(sub_boxes)
        #     obj_labels_list.append(obj_labels)
        #     obj_scores_list.append(obj_scores)
        #     sub_scores_list.append(sub_scores)
        # verb_scores = torch.cat(verb_scores_list, 1)
        # obj_boxes = torch.cat(obj_boxes_list, 1)
        # sub_boxes = torch.cat(sub_boxes_list, 1)
        # obj_labels = torch.cat(obj_labels_list, 1)
        # obj_scores = torch.cat(obj_scores_list, 1)
        # sub_scores = torch.cat(sub_scores_list, 1)
        
        obj_scores, obj_labels, obj_topk_boxes = self.get_obj_cls(outputs['pred_obj_logits'], target_sizes)
        sub_scores = self.get_sub_cls(outputs['pred_sub_logits'], target_sizes)
        obj_boxes = self.get_obj_bbox(outputs['pred_obj_boxes'], target_sizes, obj_topk_boxes)
        sub_boxes = self.get_sub_bbox(outputs['pred_sub_boxes'], target_sizes, obj_topk_boxes)
        verb_scores = self.get_verb_score(outputs['pred_verb_logits'], target_sizes, obj_topk_boxes)
    
        return sub_boxes, obj_boxes, sub_scores, obj_scores, obj_labels, verb_scores


    def get_sub_cls(self, pred_obj_logits, target_sizes):
        out_obj_logits = pred_obj_logits
        assert len(out_obj_logits) == len(target_sizes)
        obj_prob = F.softmax(out_obj_logits, -1)
        obj_scores = obj_prob[..., self.subject_category_id]
        return obj_scores


#############################################  PNMS ###########################################################
    @torch.no_grad()
    def my_triplet_nms_filter(self, preds):
        pred_od = preds['predictions']
        pred_hois = preds['hoi_prediction']
        
        # all_binary = {}     
        # for index, pred_hoi in enumerate(pred_hois):
        #     binary = str(pred_hoi['subject_id']) + '_' + \
        #              str(pred_hoi['object_id']) 
                     
        #     if binary not in all_binary: all_binary[binary] = {'scores':[], 'indexes':[]}
        #     all_binary[binary]['scores'].append(pred_hoi['score'])
        #     all_binary[binary]['indexes'].append(index)

        # for binary, values in all_binary.items():
        #     scores =  np.array(values['scores'])
        #     all_binary[binary]['std'] = scores.std()
            
        # all_binary = sorted(all_binary.values(), key=lambda k: (k.get('std', 0)), reverse=True)
        # use_binary = all_binary[:self.max_bin_num]
        
        # new_pred_hois = []
        # for _, binary in enumerate(use_binary):
        #     for _, index in enumerate(binary['indexes']):
        #         new_pred_hois.append(pred_hois[index])
        # pred_hois = new_pred_hois
        
        
        all_triplets = {}
        for index, pred_hoi in enumerate(pred_hois):
            triplet = str(pred_od[pred_hoi['subject_id']]['category_id']) + '_' + \
                      str(pred_od[pred_hoi['object_id']]['category_id']) + '_' + \
                      str(pred_hoi['category_id'])
                         
            if triplet not in all_triplets: all_triplets[triplet] = {'subs':[], 'objs':[], 'scores':[], 'std':[],'indexes':[]}
            all_triplets[triplet]['subs'].append(pred_od[pred_hoi['subject_id']]['bbox'])
            all_triplets[triplet]['objs'].append(pred_od[pred_hoi['object_id']]['bbox'])
            all_triplets[triplet]['scores'].append(pred_hoi['score'])
            all_triplets[triplet]['indexes'].append(index)
         
        all_keep_inds = []
        score_list = []
        index_list = []
        for triplet, values in all_triplets.items():
            subs, objs, scores = values['subs'], values['objs'], values['scores']
            keep_inds, scores = self.soft_pairwise_nms(np.array(subs), np.array(objs), np.array(scores))
            keep_inds = list(np.array(values['indexes'])[keep_inds])
            all_keep_inds.extend(keep_inds)
            score_list.extend(scores)
            index_list.extend(list(np.array(values['indexes'])))
            
        hoi_predictions = np.array(pred_hois)[index_list]
        score_predictions = list(np.array(score_list))
        
        for i, hoi_prediction in enumerate(hoi_predictions): 
            hoi_predictions[i]['score'] = score_predictions[i]
        hoi_predictions = list(np.array(pred_hois)[all_keep_inds])
        
        preds_filtered = {
            'predictions': pred_od,
            'hoi_prediction': hoi_predictions}
        return preds_filtered
    
    @torch.no_grad()
    def my_pairwise_nms(self, subs, objs, scores):
        sx1, sy1, sx2, sy2 = subs[:, 0], subs[:, 1], subs[:, 2], subs[:, 3]
        ox1, oy1, ox2, oy2 = objs[:, 0], objs[:, 1], objs[:, 2], objs[:, 3]
        
        sub_areas = (sx2 - sx1 + 1) * (sy2 - sy1 + 1)
        obj_areas = (ox2 - ox1 + 1) * (oy2 - oy1 + 1)

        radio_area = sub_areas/obj_areas
        
        lx, ly, hx, hy = np.minimum(sx1, ox1), np.minimum(sy1, oy1), np.maximum(sx2, ox2), np.maximum(sy2, oy2)
        union_areas = (hx - lx + 1) * (hy - ly + 1)

        hsx, hsy, hox, hoy =  (sx1 + sx2) / 2,  (sy1 + sy2) / 2,  (ox1 + ox2) / 2,  (oy1 + oy2) / 2
        hlx, hly, hhx, hhy = np.minimum(hsx, hox), np.minimum(hsy, hoy), np.maximum(hsx, hox), np.maximum(hsy, hoy)
        half_areas = (hhx - hlx + 1) * (hhy - hly + 1)
        order = scores.argsort()[::-1]
        
        keep_inds = []
        while order.size > 0:
            i = order[0]
            keep_inds.append(i)
            #####################################
            #   sub
            #####################################
            sxx1 = np.maximum(sx1[i], sx1[order[1:]])
            syy1 = np.maximum(sy1[i], sy1[order[1:]])
            sxx2 = np.minimum(sx2[i], sx2[order[1:]])
            syy2 = np.minimum(sy2[i], sy2[order[1:]])
            sw = np.maximum(0.0, sxx2 - sxx1 + 1)
            sh = np.maximum(0.0, syy2 - syy1 + 1)
            sub_inter = sw * sh
            sub_union = sub_areas[i] + sub_areas[order[1:]] - sub_inter
            #####################################
            #   obj
            #####################################
            oxx1 = np.maximum(ox1[i], ox1[order[1:]])
            oyy1 = np.maximum(oy1[i], oy1[order[1:]])
            oxx2 = np.minimum(ox2[i], ox2[order[1:]])
            oyy2 = np.minimum(oy2[i], oy2[order[1:]])
            ow = np.maximum(0.0, oxx2 - oxx1 + 1)
            oh = np.maximum(0.0, oyy2 - oyy1 + 1)
            obj_inter = ow * oh
            obj_union = obj_areas[i] + obj_areas[order[1:]] - obj_inter
            #####################################
            #   union
            #####################################
            uxx1 = np.maximum(lx[i], lx[order[1:]])
            uyy1 = np.maximum(ly[i], ly[order[1:]])
            uxx2 = np.minimum(hx[i], hx[order[1:]])
            uyy2 = np.minimum(hy[i], hy[order[1:]])
            uw = np.maximum(0.0, uxx2 - uxx1 + 1)
            uh = np.maximum(0.0, uyy2 - uyy1 + 1)
            uni_inter = uw * uh
            uni_union = union_areas[i] + union_areas[order[1:]] - uni_inter
            #####################################
            #   half_union
            #####################################
            # huxx1 = np.maximum(hlx[i], hlx[order[1:]])
            # huyy1 = np.maximum(hly[i], hly[order[1:]])
            # huxx2 = np.minimum(hhx[i], hhx[order[1:]])
            # huyy2 = np.minimum(hhy[i], hhy[order[1:]])
            # huw = np.maximum(0.0, huxx2 - huxx1 + 1)
            # huh = np.maximum(0.0, huyy2 - huyy1 + 1)
            # half_inter = huw * huh
            # half_union = half_areas[i] + half_areas[order[1:]] - half_inter

            radio = radio_area[order[1:]] / radio_area[i]
            for i, rad in enumerate(radio):
                radio[i] = 1/rad if rad>1 else rad
                

            ovr = sub_inter / sub_union  * \
                  obj_inter / obj_union  * \
                  uni_inter / uni_union  #* \
                #   radio 
                #   half_inter / half_union
                  
                  
            #   nms_alpha 1.0 nms_beta 0.5 
            # ovr = sub_inter/sub_union ** self.nms_alpha * \
            #       obj_inter/obj_union ** self.nms_beta * \
            #       uni_inter/uni_union ** self.nms_alpha * \
            #       half_inter/half_union ** self.nms_alpha            
            
            
            the = (ovr <= self.nms_thresh) #+ (radio <= 0.66) + (radio >= 1.33)
            #   保留的 
            inds = np.where(the)[0]
            order = order[inds + 1]
        return keep_inds
#############################################  PNMS ###########################################################


    @torch.no_grad()
    def soft_pairwise_nms(self, subs, objs, scores):
        sx1, sy1, sx2, sy2 = subs[:, 0], subs[:, 1], subs[:, 2], subs[:, 3]
        ox1, oy1, ox2, oy2 = objs[:, 0], objs[:, 1], objs[:, 2], objs[:, 3]
        
        sub_areas = (sx2 - sx1 + 1) * (sy2 - sy1 + 1)
        obj_areas = (ox2 - ox1 + 1) * (oy2 - oy1 + 1)

        radio_area = sub_areas/obj_areas
        
        lx, ly, hx, hy = np.minimum(sx1, ox1), np.minimum(sy1, oy1), np.maximum(sx2, ox2), np.maximum(sy2, oy2)
        union_areas = (hx - lx + 1) * (hy - ly + 1)

        hsx, hsy, hox, hoy =  (sx1 + sx2) / 2,  (sy1 + sy2) / 2,  (ox1 + ox2) / 2,  (oy1 + oy2) / 2
        hlx, hly, hhx, hhy = np.minimum(hsx, hox), np.minimum(hsy, hoy), np.maximum(hsx, hox), np.maximum(hsy, hoy)
        half_areas = (hhx - hlx + 1) * (hhy - hly + 1)
        
        order = scores.argsort()[::-1]
        order_cp = copy.deepcopy(order)
        keep_inds = []
        while order.size > 0:
            i = order[0]
            keep_inds.append(i)
            #####################################
            #   sub
            #####################################
            sxx1 = np.maximum(sx1[i], sx1[order[1:]])
            syy1 = np.maximum(sy1[i], sy1[order[1:]])
            sxx2 = np.minimum(sx2[i], sx2[order[1:]])
            syy2 = np.minimum(sy2[i], sy2[order[1:]])
            sw = np.maximum(0.0, sxx2 - sxx1 + 1)
            sh = np.maximum(0.0, syy2 - syy1 + 1)
            sub_inter = sw * sh
            sub_union = sub_areas[i] + sub_areas[order[1:]] - sub_inter
            #####################################
            #   obj
            #####################################
            oxx1 = np.maximum(ox1[i], ox1[order[1:]])
            oyy1 = np.maximum(oy1[i], oy1[order[1:]])
            oxx2 = np.minimum(ox2[i], ox2[order[1:]])
            oyy2 = np.minimum(oy2[i], oy2[order[1:]])
            ow = np.maximum(0.0, oxx2 - oxx1 + 1)
            oh = np.maximum(0.0, oyy2 - oyy1 + 1)
            obj_inter = ow * oh
            obj_union = obj_areas[i] + obj_areas[order[1:]] - obj_inter
            #####################################
            #   union
            #####################################
            uxx1 = np.maximum(lx[i], lx[order[1:]])
            uyy1 = np.maximum(ly[i], ly[order[1:]])
            uxx2 = np.minimum(hx[i], hx[order[1:]])
            uyy2 = np.minimum(hy[i], hy[order[1:]])
            uw = np.maximum(0.0, uxx2 - uxx1 + 1)
            uh = np.maximum(0.0, uyy2 - uyy1 + 1)
            uni_inter = uw * uh
            uni_union = union_areas[i] + union_areas[order[1:]] - uni_inter
            #####################################
            #   half_union
            #####################################
            # huxx1 = np.maximum(hlx[i], hlx[order[1:]])
            # huyy1 = np.maximum(hly[i], hly[order[1:]])
            # huxx2 = np.minimum(hhx[i], hhx[order[1:]])
            # huyy2 = np.minimum(hhy[i], hhy[order[1:]])
            # huw = np.maximum(0.0, huxx2 - huxx1 + 1)
            # huh = np.maximum(0.0, huyy2 - huyy1 + 1)
            # half_inter = huw * huh
            # half_union = half_areas[i] + half_areas[order[1:]] - half_inter

            radio = radio_area[order[1:]] / radio_area[i]
            for i, rad in enumerate(radio):
                radio[i] = 1/rad if rad>1 else rad
                

            ovr = sub_inter / sub_union  * \
                  obj_inter / obj_union  * \
                   uni_inter / uni_union  
                #   half_inter / half_union
                #   radio 
                #   half_inter / half_union
                           
            #   nms_alpha 1.0 nms_beta 0.5 
            # ovr = sub_inter/sub_union ** self.nms_alpha * \
            #       obj_inter/obj_union ** self.nms_beta * \
            #       uni_inter/uni_union ** self.nms_alpha * \
            #       half_inter/half_union ** self.nms_alpha            

            inds = np.where(ovr <= self.nms_thresh)[0]
            
            inds_upper = np.where(ovr > self.nms_thresh)[0]
            scores[order[inds_upper + 1]] *= 0 # ovr[inds_upper]
            
            order = order[inds + 1]
        return order_cp, scores